import gc
import multiprocessing
import time
import traceback
from datetime import datetime, timedelta

import pandas as pd
from catboost import CatBoostClassifier
from xbase_util.common_util import build_es_expression, s2date

from src.bean.circle_task_bean import CircleTaskBean
from src.bean.config_column_enum import ConfigColumnEnum
from src.bean.model_database_enum import ModelDatabaseEnum
from src.constant import predict_logger, capture_logger, expression_tag, expression
from src.extract.session import get_session
from src.model.model_common_util import get_model, print_model_list_entrance, common_predict
from src.model.v1.v1_increasing import v1_increasing_train_entrance
from src.model.v1.v1_pre_handle import v1_pre_handle
from src.util.common_util import while_input, input_is_yes, is_valid_datetime, is_int, printx
from src.util.config_manager import ConfigManager
from src.util.file_selector_util import select_multi_file


def v1_predict_file_entrance(config: ConfigManager) -> None:
    old_model = get_model(config)
    if old_model is None:
        return
    file_path_list = select_multi_file(title="[文件预测]若要用预测结果增量，请选择包含正负向样本的新数据",
                                       use_terminal=ConfigColumnEnum.use_terminal.value)
    if file_path_list is None:
        return
    model_id = old_model[ModelDatabaseEnum.id.value]
    df = pd.DataFrame()
    for p in file_path_list:
        df = pd.concat([df.reset_index(drop=True), pd.read_csv(p).reset_index(drop=True)], ignore_index=True)
    df2 = df.copy(deep=True)
    v1_pre_handle(df, model_id)
    common_predict(config, df, df2, old_model, file_path_list, increasing_method=v1_increasing_train_entrance)


def v1_start_predict_queue(predict_queue, old_model):
    model_id = old_model[ModelDatabaseEnum.id.value]
    model = CatBoostClassifier()
    model.load_model(old_model[ModelDatabaseEnum.model_path.value])
    abnormal_count = 0
    normal_count = 0
    while True:
        session, start_time, end_time = predict_queue.get()
        predict_logger.info(f"预测队列收到数据：{len(session)}")
        df = pd.DataFrame(session)
        v1_pre_handle(df, model_id)
        x_data = df[[item for item in df.columns.tolist() if item != 'is_dangerous']]
        y_pred = model.predict(x_data)
        current_abnormal_count = len([item for item in y_pred if f'{item}' == '1'])
        current_normal_count = len([item for item in y_pred if f'{item}' == '0'])
        abnormal_count += current_abnormal_count
        normal_count += current_normal_count
        predict_logger.info(
            f"实时预测结果：异常有{current_abnormal_count}个，共{abnormal_count}个｜正常有{current_normal_count}个，共{normal_count}个")


def start_capture_queue(capture_queue, predict_queue, config, arkime_exp):
    try:
        while True:
            param = capture_queue.get()
            if param is None:
                capture_logger.info("[提数队列]队列任务结束")
                break
            start_time, end_time = param
            capture_logger.info(f"[提数队列]本次时间范围：start_time:{start_time}  end_time:{end_time}")
            t = datetime.now()
            exp = build_es_expression(size=5000, start_time=start_time, end_time=end_time,
                                      arkime_expression=arkime_exp)
            session = get_session(exp, config)
            capture_logger.info(
                f"[提数队列]获取数据{len(session) if session is not None else 0}条总耗时:{datetime.now() - t}")
            if session is None or len(session) == 0:
                capture_logger.info(f"[提数队列]session数量为空")
                continue
            capture_logger.info(f"[提数队列]放入predict_queue{len(session)}条")
            predict_queue.put((session, start_time, end_time))
            capture_logger.info(f"[提数队列]队列数量:{capture_queue.qsize()}")
            time.sleep(1)
            gc.collect()
    except Exception as e:
        traceback.print_exc()
        capture_logger.info("[提数队列]获取数据队列报错")


def v1_start_get_session_queue(bean: CircleTaskBean, capture_queue):
    if bean.start_time is None:
        start_time = datetime.now() - timedelta(minutes=bean.time_before_now_m)
    else:
        start_time = s2date(bean.start_time)
    while True:
        endTime = start_time + timedelta(seconds=bean.time_span_s)
        capture_queue.put((start_time, endTime))
        time.sleep(bean.sleep_s)
        start_time = endTime


def v1_predict_realtime_entrance(config: ConfigManager) -> None:
    model_list = config.get_model_list()
    start_time = None
    time_before_now_m = None
    print_model_list_entrance(config)
    num = while_input("请输入要预测的模型id:", is_int)
    models = [model for model in model_list if f'{model[ModelDatabaseEnum.id.value]}' == f'{num}']
    if len(models) == 0:
        printx("模型不存在")
        return None
    old_model = models[0]

    if input_is_yes("是否使用固定的开始时间"):
        start_time = while_input("请输入开始时间:", is_valid_datetime)
    else:
        time_before_now_m = while_input("请输入前几分钟作为开始时间:", is_int)
    time_span_s = while_input("请输每次获取几秒的数据:", is_int)
    sleep_s = while_input("请输每次获取完成间隔几秒:", is_int)
    capture_queue = multiprocessing.Queue()
    predict_queue = multiprocessing.Queue()
    exp = expression[expression_tag]
    multiprocessing.Process(target=v1_start_predict_queue, args=(predict_queue, old_model)).start()
    multiprocessing.Process(target=start_capture_queue, args=(capture_queue, predict_queue, config, exp)).start()
    v1_start_get_session_queue(CircleTaskBean(
        start_time=start_time,
        time_before_now_m=time_before_now_m,  # 从10分钟前开始
        time_span_s=time_span_s,  # 获取10秒的数据
        sleep_s=sleep_s  # 获取完间隔10秒
    ), capture_queue)
